import torch
import torch.nn as nn


class MaskedBCELoss(nn.Module):
    def __init__(self):
        super().__init__()
        self.loss_fn = nn.BCELoss(reduction='mean')

    def forward(self, logits, label, umask):
        '''

        :param logits: [batch_size, max_len]
        :param label:  [batch_size, max_len]
        :param umask:  [batch_size, max_len]
        :return:
        '''
        mask_ = umask.eq(1)
        label = torch.masked_select(label.float(), mask_)
        logits = torch.masked_select(logits, mask_)
        loss = self.loss_fn(logits, label)
        return loss